import fitz
from PIL import Image
import pytesseract
from docx import Document
from pptx import Presentation
import io
from typing import List, Union, Optional, Dict
import logging
from PyPDF2 import PdfReader, PdfWriter
import tempfile
import os

from docx2pdf import convert as docx2pdf

from pdf2docx import Converter

from pydub import AudioSegment

from app.utils.file_utils import temp_file_manager

logger = logging.getLogger(__name__)

async def perform_ocr(content: bytes, filename: str, lang: str) -> str:
    """执行OCR识别"""
    with temp_file_manager() as tmp_mgr:
        try:
            if not content:
                raise ValueError("文件内容为空")

            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                image = Image.open(io.BytesIO(content))
                if image.mode != 'RGB':
                    image = image.convert('RGB')
                return pytesseract.image_to_string(image, lang=lang)
            else:
                # 创建临时PDF文件
                temp_pdf = tmp_mgr.create_temp_file('.pdf')
                with open(temp_pdf, 'wb') as f:
                    f.write(content)

                pdf_document = fitz.open(temp_pdf)
                if pdf_document.page_count == 0:
                    raise ValueError("PDF文件为空")

                text = ""
                for page in pdf_document:
                    pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    text += pytesseract.image_to_string(img, lang=lang) + "\n"
                return text.strip()

        except Exception as e:
            raise Exception(f"OCR处理失败: {str(e)}")

async def pdf_to_word(content: bytes) -> bytes:
    """PDF转Word"""
    try:
        doc = Document()
        pdf_document = fitz.open(stream=content, filetype="pdf")
        
        for page in pdf_document:
            text = page.get_text()
            doc.add_paragraph(text)
        
        docx_bytes = io.BytesIO()
        doc.save(docx_bytes)
        return docx_bytes.getvalue()
    except Exception as e:
        raise Exception(f"PDF转Word失败: {str(e)}")

async def pdf_to_ppt(content: bytes) -> bytes:
    """PDF转PPT"""
    try:
        prs = Presentation()
        pdf_document = fitz.open(stream=content, filetype="pdf")
        
        for page in pdf_document:
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            
            slide = prs.slides.add_slide(prs.slide_layouts[6])
            
            img_bytes = io.BytesIO()
            img.save(img_bytes, format='PNG')
            img_bytes.seek(0)
            
            slide.shapes.add_picture(
                img_bytes, 
                0, 0,
                width=prs.slide_width,
                height=prs.slide_height
            )
        
        pptx_bytes = io.BytesIO()
        prs.save(pptx_bytes)
        return pptx_bytes.getvalue()
    except Exception as e:
        raise Exception(f"PDF转PPT失败: {str(e)}")

@staticmethod
async def images_to_pdf(contents: List[bytes]) -> bytes:
    """
    将多张图片转换为PDF
    
    Args:
        contents: 图片文件内容列表
            
    Returns:
        bytes: 生成的PDF内容
    """
    with temp_file_manager() as tmp_mgr:
        try:
            # 创建新的PDF文档
            pdf = fitz.open()
            
            # 处理每张图片
            for i, content in enumerate(contents, 1):
                try:
                    # 打开并验证图片
                    image = Image.open(io.BytesIO(content))
                    
                    # 转换为RGB模式（如果需要）
                    if image.mode not in ('RGB', 'RGBA'):
                        image = image.convert('RGB')
                    
                    # 创建临时图片文件
                    img_path = tmp_mgr.create_temp_file('.png')
                    image.save(img_path, 'PNG', quality=95)
                    
                    # 将图片添加到PDF
                    pdf_page = pdf.new_page(width=image.width, height=image.height)
                    pdf_page.insert_image(pdf_page.rect, filename=img_path)
                    
                except Exception as e:
                    raise ValueError(f"处理第 {i} 张图片时失败: {str(e)}")
            
            # 保存PDF
            output = io.BytesIO()
            pdf.save(output)
            pdf.close()
            
            # 验证生成的PDF
            output.seek(0)
            try:
                check_pdf = fitz.open(stream=output.getvalue(), filetype="pdf")
                if check_pdf.page_count == 0:
                    raise ValueError("生成的PDF文件为空")
                check_pdf.close()
            except Exception as e:
                raise ValueError(f"生成的PDF文件无效: {str(e)}")
            
            return output.getvalue()
            
        except ValueError as e:
            logger.error(f"图片转PDF失败: {str(e)}")
            raise
        except Exception as e:
            logger.error(f"图片转PDF时发生未知错误: {str(e)}")
            raise ValueError("图片转PDF失败，请检查图片文件是否完整")

async def merge_pdfs(contents: List[bytes]) -> bytes:
    """
    合并多个PDF文件
    
    Args:
        contents: PDF文件内容列表
            
    Returns:
        bytes: 合并后的PDF内容
    """
    with temp_file_manager() as tmp_mgr:
        try:
            # 验证每个PDF文件
            for i, content in enumerate(contents, 1):
                # 创建临时文件来验证PDF
                temp_path = tmp_mgr.create_temp_file('.pdf')
                with open(temp_path, 'wb') as f:
                    f.write(content)
                
                try:
                    # 尝试打开PDF验证其完整性
                    pdf = fitz.open(temp_path)
                    if pdf.page_count == 0:
                        raise ValueError(f"第 {i} 个PDF文件为空")
                    pdf.close()
                except Exception as e:
                    raise ValueError(f"第 {i} 个PDF文件无效或已损坏: {str(e)}")

            # 使用PyPDF2合并PDF
            merger = PdfWriter()
            
            for content in contents:
                try:
                    pdf = PdfReader(io.BytesIO(content))
                    merger.append(pdf)
                except Exception as e:
                    raise ValueError(f"PDF合并失败: {str(e)}")
            
            # 保存合并后的文件
            output = io.BytesIO()
            merger.write(output)
            merger.close()
            
            # 验证合并后的PDF
            output.seek(0)
            try:
                pdf = fitz.open(stream=output.getvalue(), filetype="pdf")
                if pdf.page_count == 0:
                    raise ValueError("合并后的PDF文件为空")
                pdf.close()
            except Exception as e:
                raise ValueError(f"生成的PDF文件无效: {str(e)}")
            
            return output.getvalue()
            
        except ValueError as e:
            logger.error(f"PDF合并失败: {str(e)}")
            raise
        except Exception as e:
            logger.error(f"PDF合并时发生未知错误: {str(e)}")
            raise ValueError("PDF合并失败，请检查文件是否完整")

async def split_pdf(content: bytes, page_ranges: str) -> bytes:
    """拆分PDF"""
    try:
        pdf = fitz.open(stream=content, filetype="pdf")
        new_pdf = fitz.open()
        
        # 解析页码范围
        pages = []
        for part in page_ranges.split(','):
            if '-' in part:
                start, end = map(int, part.split('-'))
                pages.extend(range(start-1, end))
            else:
                page_num = int(part)
                pages.append(page_num-1)
        
        # 提取页面
        for page_num in pages:
            if 0 <= page_num < len(pdf):
                new_pdf.insert_pdf(pdf, from_page=page_num, to_page=page_num)
        
        pdf_bytes = io.BytesIO()
        new_pdf.save(pdf_bytes)
        return pdf_bytes.getvalue()
    except Exception as e:
        raise Exception(f"PDF拆分失败: {str(e)}")

class PdfService:
    SUPPORTED_FORMATS = ['docx', 'pptx']  # 添加 pptx 支持
    
    @staticmethod
    async def convert_to_format(content: bytes, target_format: str) -> bytes:
        """将PDF转换为其他格式"""
        if target_format not in PdfService.SUPPORTED_FORMATS:
            raise ValueError(f"不支持的目标格式: {target_format}")
            
        with temp_file_manager() as tmp_mgr:
            try:
                # 创建临时文件
                temp_in_path = tmp_mgr.create_temp_file('.pdf')
                with open(temp_in_path, 'wb') as f:
                    f.write(content)

                temp_out_path = tmp_mgr.create_temp_file(f'.{target_format}')

                # 根据目标格式选择转换方法
                if target_format.lower() == 'docx':
                    cv = Converter(temp_in_path)
                    cv.convert(temp_out_path)
                    cv.close()
                elif target_format.lower() == 'pptx':
                    # PDF转PPT
                    pdf_doc = fitz.open(temp_in_path)
                    prs = Presentation()
                    
                    for page_num in range(pdf_doc.page_count):
                        page = pdf_doc[page_num]
                        pix = page.get_pixmap(matrix=fitz.Matrix(3, 3), alpha=False)
                        
                        # 创建临时图片文件
                        img_path = tmp_mgr.create_temp_file('.png')
                        pix.save(img_path, "png")
                        
                        # 添加新的幻灯片
                        slide = prs.slides.add_slide(prs.slide_layouts[6])
                        
                        # 计算图片尺寸以适应幻灯片
                        img_width = prs.slide_width
                        img_height = int(img_width * pix.height / pix.width)
                        
                        if img_height > prs.slide_height:
                            img_height = prs.slide_height
                            img_width = int(img_height * pix.width / pix.height)
                        
                        # 居中放置图片
                        left = (prs.slide_width - img_width) / 2
                        top = (prs.slide_height - img_height) / 2
                        
                        # 添加图片到幻灯片
                        slide.shapes.add_picture(
                            img_path,
                            left, top,
                            width=img_width,
                            height=img_height
                        )
                    
                    # 保存PPT并关闭PDF
                    prs.save(temp_out_path)
                    pdf_doc.close()

                # 读取转换后的文件
                with open(temp_out_path, 'rb') as f:
                    result = f.read()

                return result

            except Exception as e:
                logger.error(f"转换失败: {str(e)}")
                raise

    @staticmethod
    async def convert_to_pdf(content: bytes, source_format: str) -> bytes:
        """将其他格式转换为PDF"""
        with temp_file_manager() as tmp_mgr:
            try:
                # 创建临时文件
                temp_in_path = tmp_mgr.create_temp_file(f'.{source_format}')
                with open(temp_in_path, 'wb') as f:
                    f.write(content)

                temp_out_path = tmp_mgr.create_temp_file('.pdf')

                # 根据源格式选择转换方法
                if source_format.lower() in ['doc', 'docx']:
                    docx2pdf(temp_in_path, temp_out_path)
                else:
                    raise ValueError(f"不支持的源文件格式: {source_format}")

                # 读取转换后的PDF
                with open(temp_out_path, 'rb') as f:
                    result = f.read()

                return result

            except Exception as e:
                logger.error(f"PDF转换失败: {str(e)}")
                raise

    @staticmethod
    async def split_pdf(content: bytes, start_page: int, end_page: int) -> bytes:
        """
        分割PDF文件
        
        Args:
            content: PDF文件内容
            start_page: 起始页码
            end_page: 结束页码
            
        Returns:
            bytes: 分割后的PDF内容
        """
        try:
            pdf = PdfReader(io.BytesIO(content))
            writer = PdfWriter()

            # 验证页码范围
            if start_page < 1 or end_page > len(pdf.pages) or start_page > end_page:
                raise ValueError("无效的页码范围")

            # 添加选定的页面
            for page_num in range(start_page - 1, end_page):
                writer.add_page(pdf.pages[page_num])

            output = io.BytesIO()
            writer.write(output)
            return output.getvalue()

        except Exception as e:
            logger.error(f"PDF分割失败: {str(e)}")
            raise 